﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Net.Mime;
using System.Security.Cryptography;
using System.Security.Cryptography.Pkcs;
using System.Security.Cryptography.X509Certificates;

namespace NHINDirect
{
    /// <summary>
    /// SMIME uses Pkcs (CMS) cryptography
    /// </summary>
    public class Cryptographer
    {
        public static readonly Cryptographer Default = new Cryptographer();
        
        EncryptionAlgorithm m_encryptionAlgorithm;
        DigestAlgorithm m_digestAlgorithm;
        bool m_includeEpilogue = true;

        public Cryptographer()
            : this(EncryptionAlgorithm.RSA_3DES, DigestAlgorithm.SHA1)
        {
        }

        public Cryptographer(EncryptionAlgorithm encryptionAlgorithm, DigestAlgorithm digestAlgorithm)
        {
            this.m_encryptionAlgorithm = encryptionAlgorithm;
            this.m_digestAlgorithm = digestAlgorithm;
        }

        public EncryptionAlgorithm EncryptionAlgorithm
        {
            get
            {
                return this.m_encryptionAlgorithm;
            }
            set
            {
                this.m_encryptionAlgorithm = value;
            }
        }

        public DigestAlgorithm DigestAlgorithm
        {
            get
            {
                return this.m_digestAlgorithm;
            }
            set
            {
                this.m_digestAlgorithm = value;
            }
        }
        
        /// <summary>
        /// When signing multipart messages, some mail clients do not include the multipart epilogue
        /// </summary>
        public bool IncludeMultipartEpilogueInSignature
        {
            get
            {
                return this.m_includeEpilogue;
            }            
            set
            {
                this.m_includeEpilogue = value;
            }            
        }
        
        //-----------------------------------------------------
        //
        // Encryption
        //
        //-----------------------------------------------------
        public MimeEntity Encrypt(MultipartEntity entity, X509Certificate2 encryptingCertificate)
        {
            return this.Encrypt(entity.ToEntity(), encryptingCertificate);
        }

        public MimeEntity Encrypt(MultipartEntity entity, X509Certificate2Collection encryptingCertificates)
        {
            return this.Encrypt(entity.ToEntity(), encryptingCertificates);
        }
        
        public MimeEntity Encrypt(MimeEntity entity, X509Certificate2 encryptingCertificate)
        {
            return this.Encrypt(entity, new X509Certificate2Collection(encryptingCertificate));
        }

        public MimeEntity Encrypt(MimeEntity entity, X509Certificate2Collection encryptingCertificates)
        {
            if (entity == null)
            {
                throw new ArgumentNullException();
            }

            byte[] messageBytes = EntitySerializer.Default.SerializeToBytes(entity);     // Serialize message out as ASCII encoded...
            
            byte[] encryptedBytes = this.Encrypt(messageBytes, encryptingCertificates);

            messageBytes = null;
            
            MimeEntity encryptedEntity = new MimeEntity();
            encryptedEntity.ContentType = new Header(Protocol.ContentTypeHeader, Protocol.EncryptedContentTypeHeaderValue);
            encryptedEntity.ContentTransferEncoding = new Header(Protocol.ContentTransferEncodingHeader, 
                                                                 Protocol.ToString(TransferEncoding.Base64));
                        
            encryptedEntity.Body = new Body(Convert.ToBase64String(encryptedBytes, Base64FormattingOptions.InsertLineBreaks));

            return encryptedEntity;
        }
        
        public byte[] Encrypt(byte[] content, X509Certificate2 encryptingCertificate)
        {
            return this.Encrypt(content, new X509Certificate2Collection(encryptingCertificate));
        }

        public byte[] Encrypt(byte[] content, X509Certificate2Collection encryptingCertificates)
        {
            EnvelopedCms envelope = this.CreateEncryptedEnvelope(content, encryptingCertificates);
            return envelope.Encode();
        }
                
        public EnvelopedCms CreateEncryptedEnvelope(byte[] content, X509Certificate2Collection encryptingCertificates)
        {
            if (content == null || encryptingCertificates == null || encryptingCertificates.Count == 0)
            {
                throw new ArgumentNullException();
            }

            CmsRecipientCollection recipients = new CmsRecipientCollection(SubjectIdentifierType.IssuerAndSerialNumber, encryptingCertificates);
            EnvelopedCms dataEnvelope = new EnvelopedCms(CreateDataContainer(content), this.ToAlgorithmID(this.m_encryptionAlgorithm));
            dataEnvelope.Encrypt(recipients);
            
            return dataEnvelope;
        }

        //-----------------------------------------------------
        //
        // Decryption
        //
        //-----------------------------------------------------
        //
        // Cryptography is performed only on the Mime portions of the message, not the RFC822 headers
        //
        public MimeEntity Decrypt(Message message, X509Certificate2 decryptingCertificate)
        {
            return this.Decrypt(message.ExtractMimeEntity(), decryptingCertificate);
        }
        
        public MimeEntity Decrypt(MimeEntity encryptedEntity, X509Certificate2 decryptingCertificate)
        {
            if (encryptedEntity == null || decryptingCertificate == null)
            {
                throw new ArgumentNullException();
            }
            
            if (!decryptingCertificate.HasPrivateKey)
            {
                throw new ArgumentException("Certificate has no private key");
            }
            
            encryptedEntity.VerifyMediaType(Protocol.EncryptedContentMediaType, Protocol.EncryptedContentMediaTypeAlternative);
            encryptedEntity.VerifyTransferEncoding(Protocol.TransferEncodingBase64);
            
            byte[] encryptedBytes = Convert.FromBase64String(encryptedEntity.Body.Text);                        
            byte[] decryptedBytes = this.Decrypt(encryptedBytes, decryptingCertificate);
            encryptedBytes = null;
            //
            // And turn the encrypted bytes back into an entity
            //
            return EntitySerializer.Default.Deserialize<MimeEntity>(decryptedBytes);
        }
        
        public byte[] Decrypt(byte[] encryptedContent, X509Certificate2 decryptingCertificate)
        {
            return this.Decrypt(encryptedContent, new X509Certificate2Collection(decryptingCertificate));
        }

        public byte[] Decrypt(byte[] encryptedContent, X509Certificate2Collection decryptingCertificates)
        {
            if (decryptingCertificates == null || decryptingCertificates.Count == 0)
            {
                throw new ArgumentException();
            }

            EnvelopedCms dataEnvelope = this.DeserializeEncryptionEnvelope(encryptedContent);

            dataEnvelope.Decrypt(decryptingCertificates);

            ContentInfo contentInfo = dataEnvelope.ContentInfo;
            if (contentInfo.ContentType.Value != Cryptographer.Oid_ContentType_Data.Value)
            {
                throw new NotSupportedException("Content must be of type data");
            }

            return contentInfo.Content;
        }

        public EnvelopedCms DeserializeEncryptionEnvelope(byte[] encryptedContent)
        {
            if (encryptedContent == null)
            {
                throw new ArgumentException();
            }

            EnvelopedCms dataEnvelope = new EnvelopedCms();
            dataEnvelope.Decode(encryptedContent);
            return dataEnvelope;
        }

        ContentInfo CreateDataContainer(byte[] content)
        {
            return new ContentInfo(Cryptographer.Oid_ContentType_Data, content);
        }

        //-----------------------------------------------------
        //
        // Signatures
        //
        //-----------------------------------------------------
        //
        // Cryptography is performed only on the Mime portions of the message, not the RFC822 headers
        // Some mail readers ignore the epilogue when calculating signatures!
        //
        public SignedEntity Sign(Message message, X509Certificate2 signingCertificate)
        {
            return this.Sign(message.ExtractEntityForSignature(this.m_includeEpilogue), signingCertificate);
        }
        
        public SignedEntity Sign(MimeEntity entity, X509Certificate2 signingCertificate)        
        {
            return this.Sign(entity, new X509Certificate2Collection(signingCertificate));
        }
        
        public SignedEntity Sign(MimeEntity entity, X509Certificate2Collection signingCertificates)
        {
            if (entity == null)
            {
                throw new ArgumentNullException();
            }

            byte[] entityBytes = EntitySerializer.Default.SerializeToBytes(entity);     
            MimeEntity signature = this.CreateSignatureEntity(entityBytes, signingCertificates);
                        
            return new SignedEntity(this.m_digestAlgorithm, entity, signature);
        }

        public byte[] Sign(byte[] content, X509Certificate2 signingCertificate)
        {
            SignedCms signature = this.CreateSignature(content, signingCertificate);
            return signature.Encode();
        }

        public byte[] Sign(byte[] content, X509Certificate2Collection signingCertificates)
        {
            SignedCms signature = this.CreateSignature(content, signingCertificates);
            return signature.Encode();
        }

        public MimeEntity CreateSignatureEntity(byte[] content, X509Certificate2Collection signingCertificates)
        {
            byte[] signatureBytes = this.Sign(content, signingCertificates);
            //
            // We create an entity to hold a detached signature
            //
            MimeEntity signature = new MimeEntity();
            signature.ContentType = new Header(Protocol.ContentTypeHeader, Protocol.SignatureContentTypeHeaderValue);
            signature.ContentTransferEncoding = new Header(Protocol.ContentTransferEncodingHeader,
                                                                 Protocol.ToString(TransferEncoding.Base64));

            signature.Body = new Body(Convert.ToBase64String(signatureBytes));

            return signature;
        }
        
        public SignedCms CreateSignature(byte[] content, X509Certificate2 signingCertificate)
        {
            if (content == null || signingCertificate == null)
            {
                throw new ArgumentNullException();
            }

            CmsSigner signer = this.CreateSigner(signingCertificate);
            SignedCms signature = new SignedCms(CreateDataContainer(content), true); // true: Detached Signature            
            signature.ComputeSignature(signer, true);   // true: don't prompt the user

            return signature;
        }

        public SignedCms CreateSignature(byte[] content, X509Certificate2Collection signingCertificates)
        {
            if (content == null || signingCertificates == null)
            {
                throw new ArgumentNullException();
            }

            SignedCms signature = new SignedCms(CreateDataContainer(content), true); // true: Detached Signature            
            for (int i = 0, count = signingCertificates.Count; i < count; ++i)
            {
                CmsSigner signer = this.CreateSigner(signingCertificates[i]);
                signature.ComputeSignature(signer, true);  // true: don't prompt the user
            }

            return signature;
        }

        //-----------------------------------------------------
        //
        // Signature Validation
        //
        //-----------------------------------------------------
        
        public void CheckSignature(SignedEntity signedEntity, X509Certificate2 signerCertificate)
        {
            SignedCms signatureEnvelope = DeserializeSignatureEnvelope(signedEntity);
            this.CheckSignature(signatureEnvelope.SignerInfos, signerCertificate);
        }                        
        
        public void CheckSignature(byte[] content, byte[] signatureBytes, X509Certificate2 signerCertificate)
        {
            if (signatureBytes == null || signerCertificate == null)
            {
                throw new ArgumentNullException();
            }

            SignedCms signature = this.DeserializeSignatureEnvelope(content, signatureBytes);            
            SignerInfoCollection signers = signature.SignerInfos;
            if (signers == null || signers.Count == 0)
            {
                throw new SignatureValidationException();
            }
            //
            // Find the signer
            //
            string expectedThumbprint = signerCertificate.Thumbprint;
            SignerInfo signer = signers.FindByThumbprint(signerCertificate.Thumbprint);
            if (signer == null)
            {
                throw new SignatureValidationException();            
            }

            signer.CheckSignature(true);
        }
        
        public void CheckSignature(SignerInfoCollection signers, X509Certificate2 signerCertificate)
        {
            if (signerCertificate == null)
            {
                throw new ArgumentNullException();
            }
            if (signers == null || signers.Count == 0)
            {
                throw new SignatureValidationException();
            }
            //
            // Find the signer
            //
            string expectedThumbprint = signerCertificate.Thumbprint;
            SignerInfo signer = signers.FindByThumbprint(signerCertificate.Thumbprint);
            if (signer == null)
            {
                throw new SignatureValidationException();
            }

            signer.CheckSignature(true);
        }
        
        public SignedCms DeserializeSignatureEnvelope(SignedEntity entity)
        {
            if (entity == null)
            {
                throw new SignatureValidationException();
            }

            // Serialize entity out as ASCII encoded...
            byte[] contentBytes = EntitySerializer.Default.SerializeToBytes(entity.Content);     

            byte[] signatureBytes = Convert.FromBase64String(entity.Signature.Body.Text);

            return this.DeserializeSignatureEnvelope(contentBytes, signatureBytes);
        }

        public SignedCms DeserializeSignatureEnvelope(byte[] content, byte[] signatureBytes)
        {
            SignedCms signature = new SignedCms(this.CreateDataContainer(content), true);
            signature.Decode(signatureBytes);
            
            return signature;
        }

        CmsSigner CreateSigner(X509Certificate2 cert)
        {
            CmsSigner signer = new CmsSigner(cert);

            signer.DigestAlgorithm = this.ToDigestAlgorithmOid(this.m_digestAlgorithm);
            signer.IncludeOption = X509IncludeOption.EndCertOnly;

            Pkcs9SigningTime signingTime = new Pkcs9SigningTime();
            signer.SignedAttributes.Add(signingTime);

            return signer;
        }

        //
        // OIDs
        //
        static Oid Oid_SHA1 = new Oid("1.3.14.3.2.26");
        static Oid Oid_SHA256 = new Oid("2.16.840.1.101.3.4.2.1");
        static Oid Oid_SHA512 = new Oid("2.16.840.1.101.3.4.2.3");
        static Oid Oid_RSA_ThreeDES = new Oid("1.2.840.113549.3.7");
        static Oid Oid_ContentType_Data = new Oid("1.2.840.113549.1.7.1");

        static AlgorithmIdentifier s_threeDESAlgorithmID = new AlgorithmIdentifier(Cryptographer.Oid_RSA_ThreeDES);

        Oid ToDigestAlgorithmOid(DigestAlgorithm type)
        {
            switch (type)
            {
                default:
                    throw new NotSupportedException();

                case DigestAlgorithm.SHA1:
                    return Cryptographer.Oid_SHA1;
            }
        }

        AlgorithmIdentifier ToAlgorithmID(EncryptionAlgorithm type)
        {
            switch (type)
            {
                default:
                    throw new NotSupportedException();

                case EncryptionAlgorithm.RSA_3DES:
                    return Cryptographer.s_threeDESAlgorithmID;
            }
        }
    }
        
}
